import tempfile
import os

install_warning = (
    "Note that these must be preinstalled "
    "on all nodes in the Ray cluster; it is not "
    "sufficient to install them in the runtime_env."
)

try:
    import boto3
    from smart_open import open as open_file
except ImportError:
    raise ImportError(
        "You must `pip install smart_open` and "
        "`pip install boto3` to fetch URIs in s3 "
        "bucket. " + install_warning
    )

class Boto3Storage:
    def __init__(self, addr=None, access_key=None, secret_key=None):
        self.addr = addr
        self.access_key = access_key
        self.secret_key = secret_key
        if addr is not None:
            if not addr.startswith(("http://", "https://")):
                addr = f"http://{addr}"
            os.environ["AWS_ACCESS_KEY_ID"] = access_key
            os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
            os.environ["AWS_ENDPOINT_URL"] = addr

    def put_npy(self, arr, s3_path):
        import numpy as np
        tp = {"client": boto3.client("s3")}
        with open_file(s3_path, "wb", transport_params=tp) as f:
            np.save(f, arr)

    def put_json(self, data, s3_path):
        import json
        tp = {"client": boto3.client("s3")}
        with open_file(s3_path, "w", encoding="utf-8", transport_params=tp) as f:
            data = json.dumps(data, ensure_ascii=False, indent=4)
            f.write(data)

    def put_csv(self, df, s3_path):
        tp = {"client": boto3.client("s3")}
        with open_file(s3_path, "w", encoding="utf-8", transport_params=tp) as f:
            df.to_csv(f, index=False)

    def get_npy(self, s3_path):
        import numpy as np
        tp = {"client": boto3.client("s3")}
        with open_file(s3_path, "rb", transport_params=tp) as f:
            return np.load(f)
        
    def get_json(self, s3_path):
        import json
        tp = {"client": boto3.client("s3")}
        with open_file(s3_path, "r", encoding="utf-8", transport_params=tp) as f:
            return json.load(f)

    def get_csv(self, s3_path):
        import pandas as pd
        tp = {"client": boto3.client("s3")}
        with open_file(s3_path, "r", encoding="utf-8", transport_params=tp) as f:
            return pd.read_csv(f)


